from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator
import os



'''
    precision: [T, R, K, A, M]  (iou, recall, cls, area range, max dets)
               [10, 101, 3, 4, 3]

    T 表示COCO计算时采用的10个IoU值, 从0.5到0.95每间隔0.05取一个值
    R 表示COCO计算时采用的每一个概率阈值，这里是从0到1每间隔0.01（即一个百分点）取一个值, 共101的值
    K 表示检测任务中检测的目标类别数，假设针对COCO数据集就为80
    A 表示检测任务中针对的目标尺度类型 共4个值:
        第一个表示没有限制, 
        第二个代表小目标（area < 32^2）
        第三个代表中等目标（32^2 < area < 96^2）
        第四个代表大目标（area > 96^2）
    M 表示每张图片最大检测目标个数，COCO中有1,10,100共3个值
'''

def _summarize(coco, ap=True, catId=None, iouThr=None, areaRng='all', maxDets=100):
    p = coco.params
    iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
    titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
    typeStr = '(AP)' if ap == 1 else '(AR)'
    iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
        if iouThr is None else '{:0.2f}'.format(iouThr)

    # areaRng ['all', 'small', 'medium', 'large']
    aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
    mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]

    if ap:
        # dimension of precision: [TxRxKxAxM]
        s = coco.eval['precision']
        # IoU
        if iouThr is not None:
            t = np.where(iouThr == p.iouThrs)[0]
            s = s[t]

        if isinstance(catId, int):
            s = s[:, :, catId, aind, mind]
        else:
            s = s[:, :, :, aind, mind]

    else:
        # dimension of recall: [TxKxAxM]
        s = coco.eval['recall']
        if iouThr is not None:
            t = np.where(iouThr == p.iouThrs)[0]
            s = s[t]

        if isinstance(catId, int):
            s = s[:, catId, aind, mind]
        else:
            s = s[:, :, aind, mind]

    if len(s[s > -1]) == 0:
        mean_s = -1
    else:
        mean_s = np.mean(s[s > -1])

    print_string = iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s)
    return mean_s, print_string


def cal_all_category_all_iou_ap(coco_eval):
    '''
    计算所有类别 area=all IoU 0.50:0.05:0.95 的10个分数
    '''
    ap_list = []
    iou_list = []
    for iou in np.arange(0.5, 1, 0.05):
        iou = round(iou, 2)
        mean_s, _ = _summarize(coco_eval, iouThr=iou)
        iou_list.append(iou)
        ap_list.append(mean_s)
    
    return iou_list, ap_list

def cal_every_category_all_iou_ap(coco_eval, categories_info):
    '''
    计算每一个类别 area=all IoU 0.50:0.05:0.95 的分数
    '''
    res = {}
    # 按每一类计算0.50:0.05:0.95
    for i, catId in enumerate(categories_info):
        per_ap = []
        per_iou = []
        for iou in np.arange(0.5, 1, 0.05):
            iou = round(iou, 2)
            per_iou.append(iou)
            mean_s, _ = _summarize(coco_eval, catId=i, iouThr=iou)
            per_ap.append(mean_s)
        res[catId] = {
            'iou': per_iou,
            'ap' : per_ap
        }
    
    return res

def cal_every_category_all_iou_recall(coco_eval, categories_info):
    res = {}
    # 按每一类计算0.50:0.05:0.95
    for i, catId in enumerate(categories_info):
        per_recall = []
        per_iou = []
        for iou in np.arange(0.5, 1, 0.05):
            iou = round(iou, 2)
            per_iou.append(iou)
            mean_s, _ = _summarize(coco_eval, catId=i, iouThr=iou, ap=False)
            per_recall.append(mean_s)
        res[catId] = {
            'iou': per_iou,
            'recall' : per_recall
        }
    
    return res
def cal_all_category_every_area_all_iou_ap(coco_eval, categories_info):

    area_list = ['all', 'small', 'medium', 'large']

    # 按面积大小
    res = {}
    for area in area_list:
        per_ap = []
        per_iou = []
        for iou in np.arange(0.5, 1, 0.05):
            iou = round(iou, 2)
            per_iou.append(iou)
            mean_s, _ = _summarize(coco_eval, iouThr=iou, areaRng=area)
            per_ap.append(mean_s)
        res[area] = {
            'iou': per_iou,
            'ap' : per_ap
        }
    return res

def cal_every_category_every_area_all_iou_ap(coco_eval, categories_info):
    area_list = ['all', 'small', 'medium', 'large']
    '''
        res  = {
            1:{
                'all':{
                    'iou':[],
                    'ap':[],
                }
            }
        }
    '''
    res = {}
    # 按每一类计算
    for i, catId in enumerate(categories_info):
        # 每一个面积
        res[catId] = {}
        for area in area_list:
            # 全分段
            per_ap = []
            per_iou = []
            for iou in np.arange(0.5, 1, 0.05):
                iou = round(iou, 2)
                per_iou.append(iou)
                mean_s, _ = _summarize(coco_eval, iouThr=iou, areaRng=area, catId=i)
                per_ap.append(mean_s)
            res[catId][area] = {
                'iou': per_iou,
                'ap' : per_ap
            }
    
    return res


def draw_every_category_every_area_all_iou_ap(res, filename, categories_info):

    # 图片数 == 类别数
    for cat_id in res:
        fig = plt.figure(figsize=(10,6))
        fig.suptitle(filename, fontsize=14, fontweight='bold')
        plt.xlabel('IoU', fontsize=13)
        plt.ylabel('AP', fontsize=13)
        handles_list = []
        labels_list = []
        for area in res[cat_id]:
            x = res[cat_id][area]
            h, = plt.plot(x['iou'], x['ap'], linewidth=2, marker='*', markersize=13)
            handles_list.append(h)
            labels_list.append(area)

            for a, b in zip(x['iou'], x['ap']):
                plt.text(a, b, round(b, 3), ha='left', va='bottom', fontsize=13)

        plt.legend(handles=handles_list, labels=labels_list, loc = 'best')
        plt.title('category : %s every area IoU 0.50:0.05:0.95'%(str(cat_id)), fontsize=13)
        # 剔除图框上边界和右边界的刻度
        plt.tick_params(top = 'off', right = 'off')
        x_major_locator=MultipleLocator(0.05)
        ax=plt.gca()
        ax.xaxis.set_major_locator(x_major_locator)
        plt.xlim(0.5, 1)
        plt.axhline(y=0,ls="dashdot",c="green")#添加水平直线
        plt.grid(linestyle='-.')  # 生成网格
        plt.savefig(os.path.join(filename, "类别 %s 全面积全分段表现.png"%(str(cat_id))),dpi = 600)

def draw_all_category_all_iou_ap(iou, ap, filename):
    fig = plt.figure(figsize=(10,6))
    fig.suptitle(filename, fontsize=14, fontweight='bold')
    plt.title('category=all area=all IoU 0.50:0.05:0.95', fontsize=13)
    plt.xlabel('IoU', fontsize=13)
    plt.ylabel('AP', fontsize=13)

    plt.plot(iou, ap, linewidth=2, color='steelblue', marker='*',
            markerfacecolor='r', markersize=13)

    for a, b in zip(iou, ap):
        plt.text(a, b, round(b, 3), ha='left', va='bottom', fontsize=13)

    # 剔除图框上边界和右边界的刻度
    plt.tick_params(top = 'off', right = 'off')
    x_major_locator=MultipleLocator(0.05)
    ax=plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    plt.xlim(0.5, 1)
    plt.axhline(y=0,ls="dashdot",c="green")#添加水平直线
    ap = np.array(ap)
    ap = ap[ap > -1]
    mAP = ap.sum() / 10.0
    plt.axhline(y=mAP,ls="dashdot",c="r")#添加水平直线
    plt.text(1.0, mAP, round(mAP, 3), ha='left', va='bottom', fontsize=13, c="r")
    plt.grid(linestyle='-.')  # 生成网格
    
    plt.savefig(os.path.join(filename, "所有类别在全分段表现.png"),dpi = 600)

def draw_all_category_every_area_all_iou_ap(res, filename, categories_info):
    fig = plt.figure(figsize=(10,6))
    fig.suptitle(filename, fontsize=14, fontweight='bold')
    plt.title('all category every area IoU 0.50:0.05:0.95', fontsize=13)
    plt.xlabel('IoU', fontsize=13)
    plt.ylabel('AP', fontsize=13)

    # 按类别
    handles_list = []
    labels_list = []
    for area in res:
        res_cat = res[area]
        h, = plt.plot(res_cat['iou'], res_cat['ap'], linewidth=2, marker='*', markersize=13)
        handles_list.append(h)
        labels_list.append(area)

        for a, b in zip(res_cat['iou'], res_cat['ap']):
            plt.text(a, b, round(b, 3), ha='left', va='bottom', fontsize=13)

    plt.legend(handles=handles_list, labels=labels_list, loc = 'best')

    # 剔除图框上边界和右边界的刻度
    plt.tick_params(top = 'off', right = 'off')
    x_major_locator=MultipleLocator(0.05)
    ax=plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    plt.xlim(0.5, 1)
    plt.axhline(y=0,ls="dashdot",c="green")#添加水平直线
    plt.grid(linestyle='-.')  # 生成网格

    plt.savefig(os.path.join(filename, "所有类别不同面积下全分段表现.png"),dpi = 600)

def draw_every_category_all_iou_recall(res, filename, categories_info):
    fig = plt.figure(figsize=(10,6))
    fig.suptitle(filename, fontsize=14, fontweight='bold')
    plt.title('every category area=all IoU 0.50:0.05:0.95', fontsize=13)
    plt.xlabel('IoU', fontsize=13)
    plt.ylabel('Recall', fontsize=13)

    # 按类别
    handles_list = []
    labels_list = []
    for cat_id in res:
        res_cat = res[cat_id]
        h, = plt.plot(res_cat['iou'], res_cat['recall'], linewidth=2, marker='*', markersize=13)
        handles_list.append(h)
        labels_list.append(categories_info[cat_id])

        for a, b in zip(res_cat['iou'], res_cat['recall']):
            plt.text(a, b, round(b, 3), ha='left', va='bottom', fontsize=13)

    plt.legend(handles=handles_list, labels=labels_list, loc = 'best')

    # 剔除图框上边界和右边界的刻度
    plt.tick_params(top = 'off', right = 'off')
    x_major_locator=MultipleLocator(0.05)
    ax=plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    plt.xlim(0.5, 1)
    plt.axhline(y=0,ls="dashdot",c="green")#添加水平直线
    plt.grid(linestyle='-.')  # 生成网格

    plt.savefig(os.path.join(filename, "每一个类别全面积全分段Recall表现.png"),dpi = 600)

def draw_every_category_all_iou_ap(res, filename, categories_info):
    fig = plt.figure(figsize=(10,6))
    fig.suptitle(filename, fontsize=14, fontweight='bold')
    plt.title('every category area=all IoU 0.50:0.05:0.95', fontsize=13)
    plt.xlabel('IoU', fontsize=13)
    plt.ylabel('AP', fontsize=13)

    # 按类别
    handles_list = []
    labels_list = []
    for cat_id in res:
        res_cat = res[cat_id]
        h, = plt.plot(res_cat['iou'], res_cat['ap'], linewidth=2, marker='*', markersize=13)
        handles_list.append(h)
        labels_list.append(categories_info[cat_id])

        for a, b in zip(res_cat['iou'], res_cat['ap']):
            plt.text(a, b, round(b, 3), ha='left', va='bottom', fontsize=13)

    plt.legend(handles=handles_list, labels=labels_list, loc = 'best')

    # 剔除图框上边界和右边界的刻度
    plt.tick_params(top = 'off', right = 'off')
    x_major_locator=MultipleLocator(0.05)
    ax=plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    plt.xlim(0.5, 1)
    plt.axhline(y=0,ls="dashdot",c="green")#添加水平直线
    plt.grid(linestyle='-.')  # 生成网格

    plt.savefig(os.path.join(filename, "每一个类别全面积全分段AP表现.png"),dpi = 600)



if __name__ == "__main__":
    filename = "val_cascade_rcnn_dcn_r50_vd_2xbifpn_ssld_gn_3x"
    label_json_path = './datasets/camera_0525/annotations/val.json'
    res_json_path = './bbox.json'
    # res_json_path = './filter_bbox.json'
    categories_info = {
        1: '1 (rect_eye)',
        2: '2 (sphere_eye)',
        3: '3 (box_eye)',
    }

    if not os.path.exists(filename):
        os.makedirs(filename)  

    coco_gt = COCO(label_json_path)
    coco_dt = coco_gt.loadRes(res_json_path)
    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()

    recall = cal_every_category_all_iou_recall(coco_eval, categories_info)
    #print(res)
    draw_every_category_all_iou_recall(recall, filename, categories_info)

    # 每一个类别 IoU全段
    res = cal_every_category_all_iou_ap(coco_eval, categories_info)
    draw_every_category_all_iou_ap(res, filename, categories_info)
    
    # 所有类别 IoU全段
    iou, ap = cal_all_category_all_iou_ap(coco_eval)
    draw_all_category_all_iou_ap(iou, ap, filename)

    # 所有类别 不同面积 IoU全段
    res2 = cal_all_category_every_area_all_iou_ap(coco_eval, categories_info)
    draw_all_category_every_area_all_iou_ap(res2, filename, categories_info)

    # 每一个类别 不同面积 IoU全段
    res3 = cal_every_category_every_area_all_iou_ap(coco_eval, categories_info)
    draw_every_category_every_area_all_iou_ap(res3, filename, categories_info)

    print("done")
